from transformers import BertTokenizer, BertModel
import torch
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

# 加载预训练的中文BERT模型和分词器
#tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
#model = BertModel.from_pretrained('bert-base-chinese')

# 保存到本地目录
#model.save_pretrained('./local_bert_model_cn')
#tokenizer.save_pretrained('./local_bert_model_cn')

# 离线模式下使用本地模型
tokenizer = BertTokenizer.from_pretrained('./local_bert_model_cn')
model = BertModel.from_pretrained('./local_bert_model_cn')


# 鲁迅《呐喊》中的长文本选段 - 使用更长的文本来演示分段
long_text = """
我在十八岁上离开家到外地去求学，那时是遇到相当的欢喜的。后来便在异乡漂泊多年，渐渐尝到客里悲秋的况味，变得寂寞起来。
到了二十四五岁，又回到家里来，这时候不但回乡的欢喜，也有回乡的悲哀，因为不久便觉得冷清得很。
家乡的变化很大，许多熟悉的面孔都已不见，街道也变得陌生。我常常独自一人走在曾经熟悉的路上，回忆着童年的点点滴滴。
那些年少的梦想和憧憬，如今都已随风而逝。只剩下对过去的怀念和对未来的迷茫。
"""

# 改进的长文本分段处理
def split_long_text(text, max_tokens=100):
    # 使用encode方法获取token IDs，然后解码回文本
    tokens = tokenizer.encode(text, add_special_tokens=False)
    
    segments = []
    for i in range(0, len(tokens), max_tokens):
        # 提取当前段的token IDs
        segment_tokens = tokens[i:i+max_tokens]
        # 将token IDs解码回文本
        segment_text = tokenizer.decode(segment_tokens, skip_special_tokens=True)
        segments.append(segment_text)
    
    return segments

# 将长文本分割成适合模型输入的片段
segments = split_long_text(long_text, max_tokens=50)  # 使用较小的max_tokens确保分段
print("分段后的文本片段:")
for i, seg in enumerate(segments, 1):
    print(f"段落 {i}: {seg}\n")
    print(f"段落 {i} 的token数量: {len(tokenizer.encode(seg))}\n")

# 动态生成嵌入向量
def generate_embeddings(segment, tokenizer, model):
    inputs = tokenizer(segment, return_tensors="pt", truncation=True, max_length=128, padding=True)
    with torch.no_grad():  # 添加no_grad以提高效率并减少内存使用
        outputs = model(**inputs)
    # 使用[CLS]标记的嵌入作为整个段落的表示
    embeddings = outputs.last_hidden_state[:, 0, :].detach().numpy()
    return embeddings

# 为每个段落生成嵌入
segment_embeddings = []
for segment in segments:
    embedding = generate_embeddings(segment, tokenizer, model)
    segment_embeddings.append(embedding)

print(f"成功生成了 {len(segment_embeddings)} 个段落的嵌入")

# 比较不同段落的语义相似性
if len(segment_embeddings) > 1:
    similarity_matrix = np.zeros((len(segment_embeddings), len(segment_embeddings)))
    for i in range(len(segment_embeddings)):
        for j in range(len(segment_embeddings)):
            similarity_matrix[i, j] = cosine_similarity(segment_embeddings[i], segment_embeddings[j])[0][0]
    
    print("\n段落间的语义相似性矩阵:")
    print(similarity_matrix)
    
    # 查找与特定段落最相似的段落（排除自身）
    if len(segment_embeddings) > 1:
        # 创建一个副本，将对角线设为-1以便找到最相似的非自身段落
        similarity_copy = similarity_matrix.copy()
        np.fill_diagonal(similarity_copy, -1)
        most_similar_idx = np.argmax(similarity_copy[0])
        print(f"\n与段落 1 最相似的段落是段落 {most_similar_idx + 1}")
else:
    print("文本太短，只分成了一个段落，无法计算相似性")